#include <cstdio>
#include <iostream>
#include <cnpy.h>
#include "zq_tools.h"

#include "stream_tools.h"
// #include "dwconv.h"
// #include "pwconv.h"
// #include "norm_actv.h"
// #include "maxPool.h"
// #include "bypass_unit.h"

#include "skynet_flow.h"
// #include "skynet_para.h"

int main()
{
    cnpy::npz_t feature = cnpy::npz_load("../data/feature_int32_HWC_big.npz");
    data_stream<3 * 8> s_in;

    for (unsigned i = 0; i < N_BATCH; ++i)
    {
        load_image<3, 8>(feature["big"], s_in);
    }

    data_stream<BIT_ACTV> s_in_1;
    data_stream<4 * BIT_ACTV> s_in_4;
    reduceWidth<
        3 * BIT_ACTV, 
        BIT_ACTV, 
        ROWIN * COLIN * N_BATCH
    >(s_in, s_in_1);
    expandWidth<
        BIT_ACTV, 
        N_IN * BIT_ACTV, 
        ROWIN * COLIN * 3 * N_BATCH
    >(s_in_1, s_in_4);

    bool pass = false;
    pass = check<4, 8>(s_in_4, feature["big"], N_BATCH);
    printf("input %s\n", pass ? "OK" : "Boom");
    printf("in size: %lu\n", s_in_4.size());

    constexpr unsigned OUT_WIDTH = 16;
    // constexpr unsigned OUT_SIZE = N_BATCH * ROW3 * COL3 * 10 / N_OUT;
    constexpr unsigned OUT_SIZE = N_BATCH * 14 / N_OUT;

    axiu_stream<OUT_WIDTH> s_axis_out;
    data_stream<OUT_WIDTH> s_out;
    
    // std::vector<uint32_t> buffer(N_BATCH * 20 * 40 * 768 / sizeof(uint32_t), 0);
    skynet_flow(s_in_4, s_axis_out);
    assert(s_axis_out.size() == OUT_SIZE);

    const ap_uint<(OUT_WIDTH + 7) / 8> ONES = -1;
    my_axiu<OUT_WIDTH> temp;
    for (unsigned i = 0; i < OUT_SIZE; ++i)
    {
        temp = s_axis_out.read();
        s_out.write(temp.data);
        assert(temp.keep == ONES);
        assert(temp.strb == ONES);
        if (i != OUT_SIZE - 1)
            assert(temp.last == 0x0);
        else
            assert(temp.last == 0x1);
    }

    printf("out size: %lu\n", s_out.size());
    pass = check<1, 16, int, ap_int<16>>(s_out, feature["fmax"], N_BATCH);
    printf("output %s\n", pass ? "OK" : "Boom");

    clear_stream<16>(s_out);

    // data_stream<192 * 8> s_in;
    // load_image<192, 8>(feature["conv6"], s_in);

    // data_stream<8> s_in2;
    // reduceWidth<192*8, 8, 40 * 80>(s_in, s_in2);

    // static uint64_t BP_buffer[20 * 40 * 768 / 8];
    // data_stream<8> s_reorg;
    // signal_fifo signal;
    // bypass_send_reOrg(s_in2, signal, BP_buffer);
    // bypass_recv(s_reorg, signal, BP_buffer);

    // // save_stream<1, 8>("../data/hls_debug.npy", s_reorg);

    // pass = check<1, 8>(s_reorg, feature["reorg"]);
    // printf("Bypass %s\n", pass ? "OK" : "Boom");

    return 0;
}

// int main()
// {
//     cnpy::npz_t feature = cnpy::npz_load("../data/skynet_feature_int32_HWC.npz");
//     cnpy::npz_t weight = cnpy::npz_load("../data/skynet_weight_int32_KKC_OCIC.npz");
//     ap_uint<3 * 5> L0_W[3 / 3][9];
//     ap_uint<1 * 16> L0_B[3 / 1];
//     ap_uint<1 * 16> L0_M[3 / 1];
//     ap_uint<2 * 3 * 5> L1_W[48 / 2][3 / 3];
//     ap_uint<2 * 16> L1_B[48 / 2];
//     ap_uint<2 * 16> L1_M[48 / 2];

//     bool pass = false;

//     data_stream<3 * 8> s_img;
//     load_image<3, 8>(feature["data0"], s_img);
//     pass = check<3, 8>(s_img, feature["data0"]);
//     printf("input %s\n", pass ? "OK" : "Boom");

//     data_stream<3 * 8> s_swu;
//     im2col_3x3<3 * 8, 3 * 8, 160, 320, 0x80808080>(s_img, s_swu);

//     data_stream<3 * 16> s_dwconv;
//     load_DW_weight<3, 3, 5>(weight["L0-W"], L0_W);
//     dwconv_3x3<3, 3, 8, 5, 16, 160 * 320>(s_swu, s_dwconv, L0_W);
//     // dwconv_3x3<3, 3, 8, 5, 16, 160 * 320>(s_swu, s_dwconv, L0_DW);

//     data_stream<16> s_dwconv2;
//     reduceWidth<3 * 16, 16, 160 * 320>(s_dwconv, s_dwconv2);

//     data_stream<1 * 8> s_act;
//     load_array1D<1, 3, 16>(weight["L0-B"], L0_B);
//     load_array1D<1, 3, 16>(weight["L0-M"], L0_M);
//     norm_actv<1, 3, 16, 16, 16, 20, 8, 16, 160 * 320>(s_dwconv2, s_act, L0_B, L0_M);
//     // save_stream<3, 8>("../data/hls_debug.npy", s_act);

//     pass = check<1, 8>(s_act, feature["conv1"]);
//     printf("L0 %s\n", pass ? "OK" : "Boom");

//     data_stream<3 * 8> s_act1;
//     expandWidth<8, 3 * 8, 160 * 320 * 3>(s_act, s_act1);

//     data_stream<2 * 16> s_pwconv;
//     load_PW_weight<3, 2, 3, 48, 5>(weight["L1-W"], L1_W);
//     pwconv<3, 2, 3, 48, 8, 5, 16, 160 * 320>(s_act1, s_pwconv, L1_W);

//     data_stream<2 * 8> s_act2;
//     load_array1D<2, 48, 16>(weight["L1-B"], L1_B);
//     load_array1D<2, 48, 16>(weight["L1-M"], L1_M);
//     norm_actv<2, 48, 16, 16, 16, 20, 8, 16, 160 * 320>(s_pwconv, s_act2, L1_B, L1_M);

//     pass = check<2, 8>(s_act2, feature["conv2"]);
//     printf("L1 %s\n", pass ? "OK" : "Boom");

//     data_stream<2 * 8> s_pool;
//     maxPool2x2<2, 48, 8, 160, 320>(s_act2, s_pool);

//     pass = check<2, 8>(s_pool, feature["pool1"]);
//     printf("P1 %s\n", pass ? "OK" : "Boom");

//     // clear_stream<4 * 4>(conv0_pool);
//     return 0;
// }
